{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e1bc453d-c8d3-4503-b3da-52120ad92c74",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Reflection Pattern\n",
    "\n",
    "The first pattern we are going to implement is the **reflection pattern**. \n",
    "\n",
    "---\n",
    "\n",
    "<img src=\"../img/reflection_pattern.png\" alt=\"Alt text\" width=\"600\"/>\n",
    "\n",
    "---\n",
    "\n",
    "This pattern allows the LLM to reflect and critique its outputs, following the next steps:\n",
    "\n",
    "1. The LLM **generates** a candidate output. If you look at the diagram above, it happens inside the **\"Generate\"** box.\n",
    "2. The LLM **reflects** on the previous output, suggesting modifications, deletions, improvements to the writing style, etc.\n",
    "3. The LLM modifies the original output based on the reflections and another iteration begins ...\n",
    "\n",
    "**Now, we are going to build, from scratch, each step, so that you can truly understand how this pattern works.**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7898c34d-de9a-4970-b7f4-3d86b69d45a7",
   "metadata": {},
   "source": [
    "## Generation Step"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "031f6b07-4f99-46f6-a53c-ff242585cbad",
   "metadata": {},
   "source": [
    "The first thing we need to consider is:\n",
    "\n",
    "> What do we want to generate? A poem? An essay? Python code?\n",
    "\n",
    "For this example, I've decided to test the Python coding skills of Llama3 70B (that's the LLM we are going to use for all the tutorials). In particular, we are going to ask our LLM to code a famous sorting algorithm: **Merge Sort**. \n",
    "\n",
    "---\n",
    "\n",
    "<img src=\"../img/mergesort.png\" alt=\"Alt text\" width=\"500\"/>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73f4d7b7-40bf-43b9-a626-2a11d5529ac8",
   "metadata": {},
   "source": [
    "### Groq Client and relevant imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96731d2f-a079-4e41-9756-220f02d4ebd8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from pprint import pprint\n",
    "from groq import Groq\n",
    "from dotenv import load_dotenv\n",
    "from IPython.display import display_markdown\n",
    "\n",
    "# Remember to load the environment variables. You should have the Groq API Key in there :)\n",
    "load_dotenv()\n",
    "\n",
    "client = Groq()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e644a635-e035-44e2-8c25-cee0f2b56556",
   "metadata": {},
   "source": [
    "We will start the **\"generation\"** chat history with the system prompt, as we said before. In this case, let the LLM act like a Python \n",
    "programmer eager to receive feedback / critique by the user."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12467256-c741-495a-9923-439c1fcf270d",
   "metadata": {},
   "outputs": [],
   "source": [
    "generation_chat_history = [\n",
    "    {\n",
    "        \"role\": \"system\",\n",
    "        \"content\": \"You are a Python programmer tasked with generating high quality Python code.\"\n",
    "        \"Your task is to Generate the best content possible for the user's request. If the user provides critique,\" \n",
    "        \"respond with a revised version of your previous attempt.\"\n",
    "    }\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43149b4f-54db-455f-9d39-6ad2f5c52b94",
   "metadata": {},
   "source": [
    "Now, as the user, we are going to ask the LLM to generate an implementation of the **Merge Sort** algorithm. Just add a new message with the **user** role to the chat history."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0742e7bd-4857-4ed1-a96b-37098d448bdd",
   "metadata": {},
   "outputs": [],
   "source": [
    "generation_chat_history.append(\n",
    "    {\n",
    "        \"role\": \"user\",\n",
    "        \"content\": \"Generate a Python implementation of the Merge Sort algorithm\"\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4df1bffe-375f-4a9a-8433-e217eb94aea2",
   "metadata": {},
   "source": [
    "Let's generate the first version of the essay."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff984277-733c-4495-b7fd-0669393380b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "mergesort_code = client.chat.completions.create(\n",
    "    messages=generation_chat_history,\n",
    "    model=\"llama3-70b-8192\"\n",
    ").choices[0].message.content\n",
    "\n",
    "generation_chat_history.append(\n",
    "    {\n",
    "        \"role\": \"assistant\",\n",
    "        \"content\": mergesort_code\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c03f208b-2234-4fd1-a02b-f4fff06c01a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "display_markdown(mergesort_code, raw=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a04ebe5-0573-4520-a529-aff22d486b7d",
   "metadata": {},
   "source": [
    "## Reflection Step"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67aa69e4-632f-4a0c-a6f0-c5a7ced4849d",
   "metadata": {
    "tags": []
   },
   "source": [
    "Now, let's allow the LLM to reflect on its outputs by defining another system prompt. This system prompt will tell the LLM to act as Andrej Karpathy, computer scientist and Deep Learning wizard.\n",
    "\n",
    ">To be honest, I don't think the fact of acting like Andrej Karpathy will influence the LLM outputs, but it was fun :)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "166f626f-dc32-4a2a-920c-88bee73bdc8b",
   "metadata": {},
   "source": [
    "<img src=\"../img/karpathy.png\" alt=\"Alt text\" width=\"500\"/>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d93c928-d585-48af-a74c-a5b8d84593c6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "reflection_chat_history = [\n",
    "    {\n",
    "    \"role\": \"system\",\n",
    "    \"content\": \"You are Andrej Karpathy, an experienced computer scientist. You are tasked with generating critique and recommendations for the user's code\",\n",
    "    }\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c498175f-b3f9-40af-92a3-d5b36d77d1cf",
   "metadata": {},
   "source": [
    "The user message, in this case,  is the essay generated in the previous step. We simply add the `mergesort_code` to the `reflection_chat_history`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26af1a73-4d91-40e8-a9bc-c34d32b2ab82",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "reflection_chat_history.append(\n",
    "    {\n",
    "        \"role\": \"user\",\n",
    "        \"content\": mergesort_code\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bfa994c8-3612-47b0-9571-e21d0d73d896",
   "metadata": {},
   "source": [
    "Now, let's generate a critique to the Python code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40fee42f-d47a-41b1-a40d-7208ba76ce98",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "critique = client.chat.completions.create(\n",
    "    messages=reflection_chat_history,\n",
    "    model=\"llama3-70b-8192\"\n",
    ").choices[0].message.content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0fef3203-c7f1-407f-8b9b-4e8ae140a4cb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "display_markdown(critique, raw=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5df433b0-d662-4378-895e-6b09dd3201bc",
   "metadata": {},
   "source": [
    "Finally, we just need to add this *critique* to the `generation_chat_history`, in this case, as the `user` role."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27a85bb3-cf6a-4576-8caf-cd41e602a1f1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "generation_chat_history.append(\n",
    "    {\n",
    "        \"role\": \"user\",\n",
    "        \"content\": critique\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3c1aefa-8454-41ab-af40-2675f340a577",
   "metadata": {},
   "source": [
    "## Generation Step (II)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91d845cf-51c3-4cfd-b6a7-1b970413f6db",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "essay = client.chat.completions.create(\n",
    "    messages=generation_chat_history,\n",
    "    model=\"llama3-70b-8192\"\n",
    ").choices[0].message.content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef14eaa8-f501-4efc-997f-8564ec8dccd8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "display_markdown(essay, raw=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75883af2-f31d-4c24-b1ff-315a0711f9fa",
   "metadata": {},
   "source": [
    "## And the iteration starts again ..."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5b824d1-c17e-448c-bdd7-df543aa5a9fd",
   "metadata": {},
   "source": [
    "After **Generation Step (II)** the corrected Python code will be received, once again, by Karpathy. Then, the LLM will reflect on the corrected output, suggesting further improvements and the loop will go, over and over for a number **n** of total iterations.\n",
    "\n",
    "> There's another possibility. Suppose the Reflection step can't find any further improvement. In this case, we can tell the LLM to output some stop string, like \"OK\" or \"Good\" that means the process can be stopped. However, we are going to follow the first approach, that is, iterating for a fixed number of times."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7cf2cf5b-d083-435c-914a-3ff484d53473",
   "metadata": {},
   "source": [
    "## Implementing a class "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15f9a9e6-29f3-4adf-863e-c49fbb9a6b44",
   "metadata": {},
   "source": [
    "Now that you understand the underlying loop of the Reflection Agent, let's implement this agent as a class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f904241-29a1-4519-b6ab-15be0a7cfc53",
   "metadata": {},
   "outputs": [],
   "source": [
    "from agentic_patterns import ReflectionAgent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd1a8071-c763-4dbf-8db7-60f9116f62e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "agent = ReflectionAgent()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87c8cf16-0dfa-49b6-bc30-8f14bbe7860a",
   "metadata": {},
   "outputs": [],
   "source": [
    "generation_system_prompt = \"You are a Python programmer tasked with generating high quality Python code\"\n",
    "\n",
    "reflection_system_prompt = \"You are Andrej Karpathy, an experienced computer scientist\"\n",
    "\n",
    "user_msg = \"Generate a Python implementation of the Merge Sort algorithm\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a9a3e5b-9b45-4a27-b391-f78b57ff94f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "final_response = agent.run(\n",
    "    user_msg=user_msg,\n",
    "    generation_system_prompt=generation_system_prompt,\n",
    "    reflection_system_prompt=reflection_system_prompt,\n",
    "    n_steps=10,\n",
    "    verbose=1,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b69d182-d12e-40bb-8dfb-cbc8903218a1",
   "metadata": {},
   "source": [
    "## Final result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e4663cd-61dd-4a38-866a-f032045a444a",
   "metadata": {},
   "outputs": [],
   "source": [
    "display_markdown(final_response, raw=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
